# -*- coding: utf-8 -*-
"""
处理 qa_files 的脚本：
1. 将所有没有适用地域的文件，将适用地域改为"全国"
2. 将所有没有知识类型标签的文件，知识类型改为"标准"
3. 如果文件没有要素分类，从 location 字段获取文件位置，拼成绝对路径，
   替换 .pdf 为 _results.json，读取 category 字段更新到要素分类字段上

用法示例：
  python scripts/process_qa_files.py
"""
import argparse
import json
import os
import sys
from pathlib import Path

# 兼容直接运行
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.append(str(PROJECT_ROOT))

from loguru import logger
from sqlalchemy.orm import selectinload

from db.base import SessionLocal
from db.models.files import QAFile
from db.models.attributes import ApplicableArea, KnowledgeType, ElementClassification
from db.repository.qa_file import db_file_add_asso
from db.repository.attributes import (
    db_get_all_applicable_areas,
    db_get_all_knowledge_types,
    db_get_element_classification_by_name,
)
from utils.file_utils import to_abs_path


def find_attribute_by_name(attributes, name: str):
    """根据名称查找属性，返回属性ID"""
    for attr in attributes:
        if attr.name == name:
            return attr.id
    return None


def get_or_create_element_classification(category: str):
    """根据 category 名称获取或创建要素分类"""
    if not category:
        return None
    
    # 先尝试查找
    classification = db_get_element_classification_by_name(category)
    if classification:
        return classification.id
    return None
    
    # 如果不存在，创建新的
    # try:
    #     new_classification = db_create_element_classification(category)
    #     logger.info(f"创建新的要素分类: {category} (ID: {new_classification.id})")
    #     return new_classification.id
    # except Exception as e:
    #     logger.error(f"创建要素分类失败 {category}: {e}")
    #     return None


def load_results_json(json_path: Path):
    """加载并解析 results.json 文件"""
    try:
        if not json_path.exists():
            logger.warning(f"JSON 文件不存在: {json_path}")
            return None
        
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data
    except Exception as e:
        logger.error(f"加载 JSON 文件失败 {json_path}: {e}")
        return None


def process_file(file_obj, applicable_area_id: int, knowledge_type_id: int, dry_run: bool = False):
    """
    处理单个文件
    
    Args:
        file_obj: QAFile 对象
        applicable_area_id: "全国"的适用地域ID
        knowledge_type_id: "标准"的知识类型ID
        dry_run: 是否为试运行模式（不实际更新数据库）
    
    Returns:
        (updated_fields, message) 元组
    """
    updated_fields = []
    messages = []
    
    # 1. 检查并更新适用地域
    if not file_obj.applicable_areas:
        updated_fields.append("applicable_area")
        messages.append("添加适用地域: 全国")
        if not dry_run:
            db_file_add_asso(file_obj.id, "applicable_area", [applicable_area_id])
            logger.info(f"文件 {file_obj.name} ({file_obj.id}): 添加适用地域 '全国'")
    
    # 2. 检查并更新知识类型
    if not file_obj.knowledge_types:
        updated_fields.append("knowledge_type")
        messages.append("添加知识类型: 标准")
        if not dry_run:
            db_file_add_asso(file_obj.id, "knowledge_type", [knowledge_type_id])
            logger.info(f"文件 {file_obj.name} ({file_obj.id}): 添加知识类型 '标准'")
    
    # 3. 检查并更新要素分类
    if not file_obj.element_classifications:
        if not file_obj.location:
            messages.append("跳过要素分类更新: location 字段为空")
        else:
            # 从 location 字段获取文件路径
            abs_path = to_abs_path(file_obj.location)
            
            # 替换 .pdf 为 _results.json
            if abs_path.endswith('.pdf'):
                json_path = Path(abs_path.replace('.pdf', '_results.json'))
            else:
                # 如果不是 .pdf，尝试直接替换扩展名
                json_path = Path(str(Path(abs_path).with_suffix('')) + '_results.json')
            
            # 读取 JSON 文件的 category 字段
            results_data = load_results_json(json_path)
            if results_data and 'category' in results_data:
                category = results_data['category']
                if category:
                    # 获取或创建要素分类
                    classification_id = get_or_create_element_classification(category)
                    if classification_id:
                        updated_fields.append("element_classification")
                        messages.append(f"添加要素分类: {category}")
                        if not dry_run:
                            db_file_add_asso(file_obj.id, "element_classification", [classification_id])
                            logger.info(f"文件 {file_obj.name} ({file_obj.id}): 添加要素分类 '{category}'")
                    else:
                        messages.append(f"无法获取或创建要素分类: {category}")
                else:
                    messages.append("JSON 文件中的 category 字段为空")
            else:
                messages.append(f"无法读取 JSON 文件或缺少 category 字段: {json_path}")
    
    return updated_fields, messages


def main():
    parser = argparse.ArgumentParser(
        description="处理 qa_files：更新适用地域、知识类型和要素分类"
    )
    parser.add_argument('--dry-run', action='store_true',
                       help='试运行模式，不实际更新数据库（默认: False）')
    args = parser.parse_args()
    
    # 获取所有属性，查找"全国"和"标准"的ID
    applicable_areas = db_get_all_applicable_areas()
    knowledge_types = db_get_all_knowledge_types()
    
    # 查找"全国"的ID
    applicable_area_id = find_attribute_by_name(applicable_areas, "全国")
    if applicable_area_id is None:
        logger.error("未找到 '全国' 适用地域，请检查数据库")
        sys.exit(1)
    
    # 查找"标准"的ID
    knowledge_type_id = find_attribute_by_name(knowledge_types, "标准")
    if knowledge_type_id is None:
        logger.error("未找到 '标准' 知识类型，请检查数据库")
        sys.exit(1)
    
    logger.info(f"找到 '全国' 适用地域 ID: {applicable_area_id}")
    logger.info(f"找到 '标准' 知识类型 ID: {knowledge_type_id}")
    
    # 统计信息
    stats = {
        'total': 0,
        'updated_applicable_area': 0,
        'updated_knowledge_type': 0,
        'updated_element_classification': 0,
        'skipped': 0,
        'errors': 0,
    }
    
    # 获取所有 qa_files 并处理
    session = SessionLocal()
    try:
        files = session.query(QAFile)\
            .options(selectinload(QAFile.attributes))\
            .filter(QAFile.is_delete.is_(False))\
            .all()
        
        total = len(files)
        stats['total'] = total
        logger.info(f"找到 {total} 个未删除的 qa_files")
        
        if total == 0:
            logger.warning("没有找到需要处理的文件")
            return
        
        # 处理每个文件
        for idx, file_obj in enumerate(files, start=1):
            try:
                logger.info(f"[{idx}/{total}] 处理文件: {file_obj.name} ({file_obj.id})")
                
                updated_fields, messages = process_file(
                    file_obj,
                    applicable_area_id,
                    knowledge_type_id,
                    dry_run=args.dry_run
                )
                
                if updated_fields:
                    if 'applicable_area' in updated_fields:
                        stats['updated_applicable_area'] += 1
                    if 'knowledge_type' in updated_fields:
                        stats['updated_knowledge_type'] += 1
                    if 'element_classification' in updated_fields:
                        stats['updated_element_classification'] += 1
                else:
                    stats['skipped'] += 1
                
                # 输出详细信息
                if messages:
                    for msg in messages:
                        logger.info(f"  - {msg}")
                else:
                    logger.info(f"  - 无需更新")
                    
            except Exception as e:
                stats['errors'] += 1
                logger.error(f"处理文件失败 {file_obj.name} ({file_obj.id}): {e}")
                import traceback
                logger.error(traceback.format_exc())
        
        if not args.dry_run:
            session.commit()
    except Exception as e:
        session.rollback()
        logger.error(f"处理过程中发生错误: {e}")
        import traceback
        logger.error(traceback.format_exc())
        raise
    finally:
        session.close()
    
    # 输出统计信息
    logger.info("\n" + "="*50)
    logger.info("处理完成统计:")
    logger.info(f"  总计:           {stats['total']}")
    logger.info(f"  更新适用地域:   {stats['updated_applicable_area']}")
    logger.info(f"  更新知识类型:   {stats['updated_knowledge_type']}")
    logger.info(f"  更新要素分类:   {stats['updated_element_classification']}")
    logger.info(f"  跳过:           {stats['skipped']}")
    logger.info(f"  错误:           {stats['errors']}")
    if args.dry_run:
        logger.info("  [试运行模式 - 未实际更新数据库]")
    logger.info("="*50)


if __name__ == '__main__':
    main()

